from docx import Document
import os
import tempfile
import shutil
from docx.shared import Pt, RGBColor
import chardet
from docx.enum.text import WD_ALIGN_PARAGRAPH
from style_system import WordStyleSystem
import logging

class DocFormatter:
    def __init__(self, sample_path: str = None):
        """初始化文档格式化器
        
        Args:
            sample_path: 样本文档路径，可选
        """
        self.sample_path = sample_path
        self.sample_doc = None
        self.sample_styles = {}
        self.logger = logging.getLogger(__name__)
        
        try:
            if sample_path:
                self.load_sample_styles()
            else:
                self.logger.warning("No sample document provided, using default styles")
                # 使用默认样式
                self._initialize_default_styles()
        except Exception as e:
            self.logger.error(f"Error loading sample document: {str(e)}")
            self._initialize_default_styles()
            
    def _initialize_default_styles(self):
        """初始化默认样式"""
        # 这里可以定义一些默认的样式
        self.sample_styles = {
            'Normal': {
                'font_size': 12,
                'font_name': 'Times New Roman',
                'alignment': 'left',
                'line_spacing': 1.5
            },
            'Heading 1': {
                'font_size': 16,
                'font_name': 'Times New Roman',
                'bold': True,
                'alignment': 'left',
                'line_spacing': 1.5
            },
            'Heading 2': {
                'font_size': 14,
                'font_name': 'Times New Roman',
                'bold': True,
                'alignment': 'left',
                'line_spacing': 1.5
            }
        }
    
    def load_sample_styles(self):
        """
        加载样本文档，初始化样式系统
        """
        try:
            sample_doc = Document(self.sample_path)
            self.sample_doc = sample_doc
            self.style_system = WordStyleSystem(sample_doc)
            print("\nAvailable styles:", sorted(self.style_system.style_map.keys()))
            return True
        except Exception as e:
            print(f"Error loading sample document: {e}")
            return False
    
    def format_document(self, input_path, output_path):
        """格式化单个文档"""
        try:
            # 加载输入文档
            input_doc = Document(input_path)
            
            # 创建新文档
            new_doc = Document()
            
            # 复制样本文档的所有样式
            if not os.path.exists(self.sample_path):
                raise FileNotFoundError(f"Sample document not found at: {self.sample_path}")
            sample_doc = Document(self.sample_path)
            
            # 使用样式系统初始化样式
            style_system = WordStyleSystem(sample_doc)
            
            # 先复制所有样式到新文档
            for style in sample_doc.styles:
                try:
                    # 创建新样式
                    if style.name not in new_doc.styles:
                        new_style = new_doc.styles.add_style(style.name, style.type)
                        
                        # 设置样式属性
                        if style.type == 'paragraph':
                            # 设置段落格式
                            if style.paragraph_format:
                                new_style.paragraph_format.alignment = style.paragraph_format.alignment
                                new_style.paragraph_format.space_before = style.paragraph_format.space_before
                                new_style.paragraph_format.space_after = style.paragraph_format.space_after
                                new_style.paragraph_format.line_spacing = style.paragraph_format.line_spacing
                                new_style.paragraph_format.first_line_indent = style.paragraph_format.first_line_indent
                                new_style.paragraph_format.left_indent = style.paragraph_format.left_indent
                                new_style.paragraph_format.right_indent = style.paragraph_format.right_indent
                            
                            # 设置字体属性
                            if style.font:
                                new_style.font.name = style.font.name
                                new_style.font.size = style.font.size
                                new_style.font.bold = style.font.bold
                                new_style.font.italic = style.font.italic
                                new_style.font.underline = style.font.underline
                                if style.font.color:
                                    new_style.font.color.rgb = style.font.color.rgb
                                if style.font.highlight_color:
                                    new_style.font.highlight_color = style.font.highlight_color
                        
                        # 设置基础样式
                        if style.base_style:
                            if style.base_style.name in new_doc.styles:
                                new_style.base_style = new_doc.styles[style.base_style.name]
                            else:
                                # 如果基础样式不存在，使用 Normal
                                if style.base_style.name.lower().startswith('normal'):
                                    new_style.base_style = new_doc.styles['Normal']
                                else:
                                    print(f"Warning: Base style '{style.base_style.name}' not found in new document")
                except Exception as e:
                    print(f"Warning: Failed to copy style {style.name}: {e}")
                    continue
            
            # 处理每个段落
            for para in input_doc.paragraphs:
                try:
                    # 跳过空段落
                    if not para.text.strip():
                        continue
                        
                    # 创建新段落
                    new_para = new_doc.add_paragraph()
                    
                    # 应用段落样式
                    if para.style and para.style.name:
                        try:
                            # 获取样本文档中的原始样式
                            original_style = sample_doc.styles[para.style.name]
                            
                            # 在新文档中创建或更新样式
                            if para.style.name not in new_doc.styles:
                                new_style = new_doc.styles.add_style(para.style.name, original_style.type)
                            else:
                                new_style = new_doc.styles[para.style.name]
                            
                            # 复制所有样式属性
                            new_style.font.name = original_style.font.name
                            new_style.font.size = original_style.font.size
                            new_style.font.bold = original_style.font.bold
                            new_style.font.italic = original_style.font.italic
                            new_style.font.underline = original_style.font.underline
                            if original_style.font.color:
                                new_style.font.color.rgb = original_style.font.color.rgb
                            if original_style.font.highlight_color:
                                new_style.font.highlight_color = original_style.font.highlight_color
                            
                            # 复制段落格式
                            if original_style.type == 'paragraph':
                                if hasattr(original_style, 'paragraph_format'):
                                    # 设置对齐方式
                                    if hasattr(original_style.paragraph_format, 'alignment'):
                                        new_style.paragraph_format.alignment = original_style.paragraph_format.alignment
                                    
                                    # 设置间距
                                    if hasattr(original_style.paragraph_format, 'space_before'):
                                        new_style.paragraph_format.space_before = original_style.paragraph_format.space_before
                                    if hasattr(original_style.paragraph_format, 'space_after'):
                                        new_style.paragraph_format.space_after = original_style.paragraph_format.space_after
                                    
                                    # 设置行间距
                                    if hasattr(original_style.paragraph_format, 'line_spacing'):
                                        new_style.paragraph_format.line_spacing = original_style.paragraph_format.line_spacing
                                    if hasattr(original_style.paragraph_format, 'line_spacing_rule'):
                                        new_style.paragraph_format.line_spacing_rule = original_style.paragraph_format.line_spacing_rule
                                    
                                    # 处理列表格式
                                    if hasattr(original_style.paragraph_format, '_element'):
                                        # 获取列表级别
                                        numPr = original_style.paragraph_format._element.xpath('w:numPr')
                                        if numPr:
                                            ilvl = numPr[0].xpath('w:ilvl')
                                            if ilvl:
                                                list_level = int(ilvl[0].get('{http://schemas.openxmlformats.org/wordprocessingml/2006/main}val'))
                                                new_style.paragraph_format.left_indent = list_level * 36  # 36磅为一个缩进级别
                                            
                                            # 获取列表样式
                                            numId = numPr[0].xpath('w:numId')
                                            if numId:
                                                list_style = int(numId[0].get('{http://schemas.openxmlformats.org/wordprocessingml/2006/main}val'))
                                                # 设置列表样式
                                                if hasattr(new_style.paragraph_format, '_element'):
                                                    new_style.paragraph_format._element.append(
                                                        numPr[0].clone()
                                                    )
                                    
                                    # 处理缩进
                                    if hasattr(original_style.paragraph_format, 'left_indent'):
                                        # 如果不是列表格式，使用原始缩进
                                        if not hasattr(original_style.paragraph_format, '_element'):
                                            new_style.paragraph_format.left_indent = original_style.paragraph_format.left_indent
                                    if hasattr(original_style.paragraph_format, 'right_indent'):
                                        new_style.paragraph_format.right_indent = original_style.paragraph_format.right_indent
                                    
                                    # 处理首行缩进和悬挂缩进
                                    if hasattr(original_style.paragraph_format, 'first_line_indent'):
                                        # 如果是悬挂缩进（负值）
                                        if original_style.paragraph_format.first_line_indent < 0:
                                            new_style.paragraph_format.first_line_indent = 0
                                            new_style.paragraph_format.left_indent = abs(original_style.paragraph_format.first_line_indent)
                                        # 如果是首行缩进（正值）
                                        else:
                                            new_style.paragraph_format.first_line_indent = original_style.paragraph_format.first_line_indent
                                            # 如果没有设置左缩进，则使用首行缩进值
                                            if not hasattr(original_style.paragraph_format, 'left_indent'):
                                                new_style.paragraph_format.left_indent = original_style.paragraph_format.first_line_indent
                                    
                                    # 设置其他格式
                                    if hasattr(original_style.paragraph_format, 'keep_together'):
                                        new_style.paragraph_format.keep_together = original_style.paragraph_format.keep_together
                                    if hasattr(original_style.paragraph_format, 'keep_with_next'):
                                        new_style.paragraph_format.keep_with_next = original_style.paragraph_format.keep_with_next
                                    if hasattr(original_style.paragraph_format, 'page_break_before'):
                                        new_style.paragraph_format.page_break_before = original_style.paragraph_format.page_break_before
                                    if hasattr(original_style.paragraph_format, 'widow_control'):
                                        new_style.paragraph_format.widow_control = original_style.paragraph_format.widow_control
                                    
                                    # 复制制表位
                                    if hasattr(original_style.paragraph_format, 'tab_stops'):
                                        for tab in original_style.paragraph_format.tab_stops:
                                            new_style.paragraph_format.tab_stops.add_tab_stop(
                                                tab.position,
                                                alignment=tab.alignment,
                                                leader=tab.leader
                                            )
                            
                            # 设置基础样式
                            if original_style.base_style:
                                if original_style.base_style.name in new_doc.styles:
                                    new_style.base_style = new_doc.styles[original_style.base_style.name]
                                else:
                                    print(f"Warning: Base style '{original_style.base_style.name}' not found in new document")
                            
                            # 应用样式到段落
                            new_para.style = para.style.name
                            
                            # 使用样式系统应用样式
                            style_system.apply_style(new_para, para.style.name)
                        except KeyError:
                            print(f"Warning: Style '{para.style.name}' not found in sample document. Using Normal style.")
                            new_para.style = 'Normal'
                            style_system.apply_style(new_para, 'Normal')
                    else:
                        new_para.style = 'Normal'
                        style_system.apply_style(new_para, 'Normal')
                    
                    # 处理文本运行
                    for run in para.runs:
                        if not run.text.strip():
                            continue
                            
                        # 创建新的运行
                        new_run = new_para.add_run(run.text)
                        
                        # 应用运行样式
                        if run.style and run.style.name:
                            try:
                                # 先设置运行样式
                                if run.style.name in new_doc.styles:
                                    new_run.style = run.style.name
                                else:
                                    # 如果样式不存在，使用默认样式
                                    print(f"Warning: Run style '{run.style.name}' not found in new document. Using Default Paragraph Font.")
                                    new_run.style = 'Default Paragraph Font'
                                
                                # 使用样式系统应用样式
                                style_system.apply_style(new_run, run.style.name)
                            except KeyError:
                                print(f"Warning: Run style '{run.style.name}' not found in sample document. Using Default Paragraph Font.")
                                new_run.style = 'Default Paragraph Font'
                                style_system.apply_style(new_run, 'Default Paragraph Font')
                except Exception as e:
                    print(f"Warning: Error processing paragraph: {e}")
                    continue
            
            # 保存文档
            new_doc.save(output_path)
            print(f"Document formatted and saved to: {output_path}")
            return True
        except Exception as e:
            print(f"Error formatting document: {e}")
            return False
    
    def batch_format(self, input_dir, output_dir):
        """批量格式化文档"""
        try:
            # 创建输出目录
            if not os.path.exists(output_dir):
                os.makedirs(output_dir)
            
            # 获取所有.docx文件
            docx_files = [f for f in os.listdir(input_dir) if f.endswith('.docx')]
            
            # 处理每个文件
            for docx_file in docx_files:
                input_path = os.path.join(input_dir, docx_file)
                output_path = os.path.join(output_dir, f"formatted_{docx_file}")
                
                print(f"\nProcessing: {docx_file}")
                if self.format_document(input_path, output_path):
                    print(f"Successfully processed: {docx_file}")
                else:
                    print(f"Failed to process: {docx_file}")
            
            print("\nProcessing complete!")
            print(f"Successfully formatted {len(docx_files)} files")
            print(f"Output files saved to: {output_dir}")
            
        except Exception as e:
            print(f"Error in batch processing: {e}")

if __name__ == "__main__":
    formatter = DocFormatter("sample.docx")
    formatter.batch_format("input_docs", "output_docs")